{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building a RAG Pipeline over Legal Documents\n",
    "\n",
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_cloud_services/blob/main/examples/parse/multimodal/legal_rag.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
    "\n",
    "This example shows how LlamaParse and LlamaIndex can be used to parse various types of legal documents, which may contain complex tabular data. The advantage of this is being able to quickly retrieve a specific answer to a legal question with comprehensive context — knowledge of precedents, statutes, and cases presented in the given documents. A user can quickly find the answer to or find out more details about a specific legal question without having to read through the often long documents by using LLMs.\n",
    "\n",
    "In this example, we will be using legal documents from the archive of the Library of Congress ([link to dataset](https://www.loc.gov/item/2020445568/)). These documents vary by format, with some containing pure text and others containing headings, sections, and large tables. This shows how LlamaParse can parse a wide variety of documents and still retrieve accurate results.\n",
    "\n",
    "The documents in this example include:\n",
    "- [APA Program Report](https://www.irs.gov/pub/irs-apa/a_2003-19.pdf)\n",
    "- [2004 Report on the CRA performance of Barre Savings Bank in Barre, MA](https://github.com/user-attachments/files/16536412/barre_savings_bank_evaluation.pdf)\n",
    "- [2016 Energy Supply/Demand Forecast](https://github.com/user-attachments/files/16536415/energy_supply_demand.pdf)\n",
    "- [Transcript of Senate Committee Hearing about Foreign Markets](https://github.com/user-attachments/files/16536422/foreign_markets.pdf)\n",
    "- [A Motion To Stay for an Indiana Court Case](https://github.com/user-attachments/files/16536427/motion_to_stay.pdf)\n",
    "- [Article About an OC Representative's Bill to Introduce Offshore Drilling to CA](https://github.com/user-attachments/files/16536437/oc_bill_offshore_drilling.pdf)\n",
    "- [Charter of the Subcommittee on Ocean Science and Technology](https://github.com/user-attachments/files/16536445/ost_subcommittee_charter.pdf)\n",
    "- [US Immigration Case](https://github.com/user-attachments/files/16536446/us_immigration_case.pdf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup and Installation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Install LlamaIndex, download the data, and set your API keys."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install \"llama-index>=0.13.0<0.14.0\" llama-cloud-services"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://github.com/user-attachments/files/16447759/data.zip -O data.zip\n",
    "!unzip -o data.zip\n",
    "!rm data.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set up your OpenAI and LlamaCloud keys."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"sk-...\"\n",
    "os.environ[\"LLAMA_CLOUD_API_KEY\"] = \"llx-...\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Code Implementation\n",
    "\n",
    "Set up LlamaParse. We want to parse the PDF files into markdown, translating the tabular data into markdown tables. To ensure accuracy, we will use the GPT-4o multimodal model to parse the PDFs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_cloud_services import LlamaParse\n",
    "\n",
    "parser = LlamaParse(\n",
    "    parse_mode=\"parse_page_with_agent\",\n",
    "    model=\"openai-gpt-4-1-mini\",\n",
    "    high_res_ocr=True,\n",
    "    adaptive_long_table=True,\n",
    "    outlined_table_extraction=True,\n",
    "    output_tables_as_HTML=True,\n",
    ")\n",
    "\n",
    "DATA_DIR = \"data\"\n",
    "\n",
    "\n",
    "def get_data_files(data_dir=DATA_DIR) -> list[str]:\n",
    "    files = []\n",
    "    for f in os.listdir(data_dir):\n",
    "        fname = os.path.join(data_dir, f)\n",
    "        if os.path.isfile(fname):\n",
    "            files.append(fname)\n",
    "    return files\n",
    "\n",
    "\n",
    "files = get_data_files()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load data from parser into documents containing parsed Markdown text from the legal document PDFs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Getting job results:   0%|          | 0/8 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Started parsing the file under job_id dad7b215-360c-46a6-857e-983249441395\n",
      "Started parsing the file under job_id bcfb24fb-0b30-4bd5-a87d-2a81b2d4298a\n",
      "Started parsing the file under job_id 50417384-e3fa-44fa-9f58-8344c129cedf\n",
      "Started parsing the file under job_id 49b0620f-e9fa-4736-801f-aadd6d6e21dd\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Getting job results:  12%|█▎        | 1/8 [00:23<02:43, 23.42s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Started parsing the file under job_id 729ceca5-2940-406d-b29a-0252dbf11e15\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Getting job results:  38%|███▊      | 3/8 [00:41<00:56, 11.20s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Started parsing the file under job_id 0733a9c5-d4a6-4242-9bd2-f61e931424dd\n",
      "Started parsing the file under job_id a948a2f8-521a-412a-9cbd-4574814a8d2c\n",
      "."
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Getting job results:  50%|█████     | 4/8 [00:44<00:32,  8.19s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Started parsing the file under job_id d9929a63-4f84-4567-abd9-bc352eee1db0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Getting job results:  75%|███████▌  | 6/8 [01:07<00:19,  9.70s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "...."
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Getting job results:  88%|████████▊ | 7/8 [02:47<00:39, 39.42s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "."
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Getting job results: 100%|██████████| 8/8 [03:32<00:00, 26.61s/it]\n"
     ]
    }
   ],
   "source": [
    "results = await parser.aparse(files)\n",
    "\n",
    "documents = []\n",
    "for result in results:\n",
    "    documents.extend(result.get_markdown_documents(split_by_page=True))\n",
    "\n",
    "for document in documents:\n",
    "    document.metadata[\n",
    "        \"context\"\n",
    "    ] = \"US legal documents provided by the Library of Congress.\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Setup LlamaIndex for querying the data using RAG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import (\n",
    "    VectorStoreIndex,\n",
    "    Settings,\n",
    ")\n",
    "from llama_index.embeddings.openai import OpenAIEmbedding\n",
    "from llama_index.llms.openai import OpenAI\n",
    "\n",
    "embed_model = OpenAIEmbedding(model=\"text-embedding-3-large\")\n",
    "llm = OpenAI(\"gpt-5-mini\")\n",
    "\n",
    "Settings.llm = llm\n",
    "Settings.embed_model = embed_model\n",
    "\n",
    "index = VectorStoreIndex.from_documents(documents)\n",
    "\n",
    "query_engine = index.as_query_engine()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example Queries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "The majority went to residential real estate lending—primarily 1–4 family mortgages (about 78.7% of loans, with home equity lines adding another 8.7%, for a total of 87.4%)."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.display import display, Markdown\n",
    "\n",
    "response = query_engine.query(\n",
    "    \"Where did the majority of Barre Savings Bank's loans go?\"\n",
    ")\n",
    "display(Markdown(str(response)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "He says foreign markets (especially U.S. markets) are attractive because:\n",
       "- The underlying economic data are strong.\n",
       "- The high‑tech boom creates a “halo” that attracts attention and investment.\n",
       "- There is broad, nearly bipartisan political/economic stability.\n",
       "- Foreign investors expect sensible foreign‑policy behavior and a firm monetary policy response to any rise in inflation.\n",
       "- Large foreign institutions (investment funds, insurers, banks) therefore see the markets as a safe, desirable place to put money."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "response = query_engine.query(\n",
    "    \"Why does Mr. Kubarych believe foreign markets are so important?\"\n",
    ")\n",
    "display(Markdown(str(response)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "House Democrats — including Speaker Nancy Pelosi — and other Democratic lawmakers oppose drilling off the California coast. They say it should be \"off the table\" for environmental reasons, point out that there are already millions of acres of federal lands leased to oil companies that aren’t being developed, and note oil companies have told Pelosi those leases aren’t being used because development takes a long time or the companies lack the equipment. No Democrats signed on to the proposed bill."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "response = query_engine.query(\n",
    "    \"Who is against the proposal of offshore drilling in CA and why?\"\n",
    ")\n",
    "display(Markdown(str(response)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "To advise and assist the Committee on Environment, Natural Resources, and Sustainability on national ocean science and technology issues and to advance federal ocean S&T goals by developing coordinated interagency strategies. It also serves as the National Ocean Council’s Ocean Science and Technology Interagency Policy Committee and retains the mandated functions of the prior joint subcommittee. Key roles include fostering national ocean S&T priorities; facilitating interagency coordination of research, technology, infrastructure, education, and observation/mapping programs; expanding fundamental knowledge of the ocean and its links to the Earth system and society; advancing modeling and forecasting; advising on science and technology for ecosystem-based management and stewardship; supporting use of ocean S&T in coastal and marine policy; and recommending scientific and technical assessments."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "response = query_engine.query(\n",
    "    \"What is the purpose of the Ocean Science and Technology Subcommittee?\"\n",
    ")\n",
    "display(Markdown(str(response)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "The appeal was dismissed because the petitioner is not a U.S. citizen, and the K‑1 fiancé(e) classification (Form I‑129F) is available only for fiancés of U.S. citizens. The denial is without prejudice, so the petitioner may file a new I‑129F if he becomes a U.S. citizen."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "response = query_engine.query(\"Why is the immigration appeal dismissed?\")\n",
    "display(Markdown(str(response)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "An advance pricing agreement (APA) is a binding contract between a taxpayer and the IRS that establishes an approved transfer pricing method (TPM) for specified related‑party (covered) transactions and tax years. If the taxpayer files its returns consistent with the agreed TPM, the IRS agrees not to seek an adjustment under IRC § 482 for those transactions. An APA can be unilateral (between the taxpayer and the IRS) or bilateral/multilateral (also agreeing with one or more foreign competent authorities), and is intended to resolve transfer‑pricing disputes in advance and, where bilateral, to reduce the risk of double taxation."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "response = query_engine.query(\"What is an advance pricing agreement?\")\n",
    "display(Markdown(str(response)))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
